#!/usr/bin/env python3
"""
快速Kafka测试数据生成脚本
简化版本，快速生成测试数据
"""

import json
import random
import time
from datetime import datetime, timedelta
from kafka import KafkaProducer
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
logger = logging.getLogger(__name__)

def create_producer():
    """创建Kafka生产者"""
    return KafkaProducer(
        bootstrap_servers=['192.168.37.11:9092'],
        value_serializer=lambda v: json.dumps(v, ensure_ascii=False).encode('utf-8'),
        key_serializer=lambda k: str(k).encode('utf-8') if k else None,
        acks=1,
        batch_size=16384,
        linger_ms=10,
        compression_type='gzip'
    )

def generate_message(msg_id):
    """生成单条消息"""
    # 消息模板
    templates = [
        "密码将在 {} 天后过期",
        "API 调用频率超限", 
        "备份任务执行成功",
        "新用户注册: user_{}",
        "服务器CPU使用率达到 {}%",
        "订单 ORD_{} 已完成支付",
        "系统维护将在 {} 小时后开始",
        "数据库连接异常",
        "文件上传完成: file_{}.txt",
        "邮件发送成功: {} 封"
    ]
    
    # 随机选择模板和参数
    template = random.choice(templates)
    
    if "{}" in template:
        if "天" in template:
            content = template.format(random.randint(1, 30))
        elif "%" in template:
            content = template.format(random.randint(70, 99))
        elif "用户" in template:
            content = template.format(random.randint(1000, 9999))
        elif "订单" in template:
            content = template.format(random.randint(100000, 999999))
        elif "小时" in template:
            content = template.format(random.randint(1, 24))
        elif "文件" in template:
            content = template.format(random.randint(1, 1000))
        elif "封" in template:
            content = template.format(random.randint(10, 500))
        else:
            content = template.format(random.randint(1, 100))
    else:
        content = template
    
    # 生成时间戳 (最近7天内)
    base_time = datetime.now()
    random_offset = timedelta(
        days=random.randint(0, 7),
        hours=random.randint(0, 23),
        minutes=random.randint(0, 59),
        seconds=random.randint(0, 59)
    )
    message_time = base_time - random_offset
    
    # 构建消息
    message = {
        "id": msg_id,
        "type": random.choice(["message", "notification", "alert"]),
        "content": content,
        "timestamp": message_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "status": random.choice(["已送达", "已读", "已发送", "处理中"])
    }
    
    # 随机添加额外字段
    if random.random() < 0.3:
        message["priority"] = random.choice(["low", "medium", "high"])
    
    if random.random() < 0.2:
        message["user_id"] = f"user_{random.randint(1000, 9999)}"
    
    return message

def main():
    """主函数"""
    TOPIC_NAME = "test_kafka"
    TOTAL_MESSAGES = 100000
    BATCH_SIZE = 1000
    
    logger.info(f"开始生成 {TOTAL_MESSAGES} 条测试消息到主题: {TOPIC_NAME}")
    
    try:
        # 创建生产者
        producer = create_producer()
        logger.info("Kafka生产者创建成功")
        
        start_time = time.time()
        sent_count = 0
        
        # 分批发送
        for batch_start in range(1, TOTAL_MESSAGES + 1, BATCH_SIZE):
            batch_end = min(batch_start + BATCH_SIZE - 1, TOTAL_MESSAGES)
            
            # 发送当前批次
            for msg_id in range(batch_start, batch_end + 1):
                message = generate_message(msg_id)
                key = f"key_{msg_id}"
                
                try:
                    producer.send(TOPIC_NAME, value=message, key=key)
                    sent_count += 1
                except Exception as e:
                    logger.error(f"发送消息 {msg_id} 失败: {e}")
            
            # 每批次后刷新并显示进度
            producer.flush()
            progress = (batch_end / TOTAL_MESSAGES) * 100
            logger.info(f"进度: {progress:.1f}% ({batch_end}/{TOTAL_MESSAGES})")
            
            # 短暂休息，避免过载
            time.sleep(0.05)
        
        # 最终刷新
        producer.flush()
        
        # 统计结果
        end_time = time.time()
        duration = end_time - start_time
        
        logger.info("=" * 50)
        logger.info("✅ 数据生成完成!")
        logger.info(f"📊 总消息数: {TOTAL_MESSAGES}")
        logger.info(f"✅ 成功发送: {sent_count}")
        logger.info(f"⏱️  耗时: {duration:.2f} 秒")
        logger.info(f"🚀 平均速度: {sent_count/duration:.0f} 消息/秒")
        logger.info("=" * 50)
        
    except KeyboardInterrupt:
        logger.info("❌ 用户中断操作")
    except Exception as e:
        logger.error(f"❌ 生成数据失败: {e}")
    finally:
        if 'producer' in locals():
            producer.close()
            logger.info("🔒 Kafka生产者已关闭")

if __name__ == "__main__":
    main()
